package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo4TranMode {

  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    //1、读取处理好的数据
    val imageData: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", "784") //知道特征数量
      .load("spark/data/images")

    //2、将数据拆分成训练集测试集
    //训练集用于训练模型
    //测试用于测试模型的准确率
    val Array(tran: DataFrame, test: DataFrame) = imageData.randomSplit(Array(7, 3))

    //3、选择算法
    //数字识别是一个分类问题，可以选择逻辑回归
    val lr: LogisticRegression = new LogisticRegression()
      //.setMaxIter(10) //最大迭代次数

    //4、将训练集带入算法训练模型
    val model: LogisticRegressionModel = lr.fit(tran)

    //5、使用测试数据测试模型的准确率
    val frame: DataFrame = model.transform(test)

    //6、计算准确率
    val p: Double = frame.where($"label" === $"prediction").count() / frame.count().toDouble


    println(s"模型的准确率：${p}")


    //7、保存模型
    model.save("spark/data/image_model")
  }

}
